{
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.14",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [],
   "dockerImageVersionId": 30776,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook",
   "isGpuEnabled": true
  }
 },
 "nbformat_minor": 5,
 "nbformat": 4,
 "cells": [
  {
   "cell_type": "code",
   "source": "import torch\nfrom collections import OrderedDict\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport scipy.io\n#from pyDOE import lhs#拉丁超立方抽样\nfrom torch import nn\nimport time\nimport warnings\nwarnings.filterwarnings(\"ignore\", category=UserWarning)  # 只忽略UserWarning类型的警告",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:44.161933Z",
     "iopub.execute_input": "2024-09-24T08:23:44.162356Z",
     "iopub.status.idle": "2024-09-24T08:23:45.981338Z",
     "shell.execute_reply.started": "2024-09-24T08:23:44.162314Z",
     "shell.execute_reply": "2024-09-24T08:23:45.980437Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T08:49:27.326666Z",
     "start_time": "2024-10-06T08:49:25.758416Z"
    }
   },
   "id": "af4b048c0e49992d",
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "if torch.cuda.is_available():\n    device = torch.device('cuda')\nelse:\n    device = torch.device('cpu')\ndef setup_seed(seed):\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    np.random.seed(seed)\n    #  random.seed(seed)\n    torch.backends.cudnn.deterministic = True\n\nsetup_seed(124)\ndtype = torch.double",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:45.986195Z",
     "iopub.execute_input": "2024-09-24T08:23:45.988477Z",
     "iopub.status.idle": "2024-09-24T08:23:46.049102Z",
     "shell.execute_reply.started": "2024-09-24T08:23:45.988430Z",
     "shell.execute_reply": "2024-09-24T08:23:46.048271Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T08:49:27.341961Z",
     "start_time": "2024-10-06T08:49:27.327668Z"
    }
   },
   "id": "9e2f9d35aa33fbb9",
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "#求导\ndef gradients(outputs, inputs):\n    return torch.autograd.grad(outputs, inputs, grad_outputs=torch.ones_like(outputs), create_graph=True)\n#类型转换\ndef to_numpy(input):\n    if isinstance(input, torch.Tensor):\n        return input.detach().cpu().numpy()\n    elif isinstance(input, np.ndarray):\n        return input\n",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:46.052132Z",
     "iopub.execute_input": "2024-09-24T08:23:46.052502Z",
     "iopub.status.idle": "2024-09-24T08:23:46.057747Z",
     "shell.execute_reply.started": "2024-09-24T08:23:46.052456Z",
     "shell.execute_reply": "2024-09-24T08:23:46.056804Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T08:49:27.357238Z",
     "start_time": "2024-10-06T08:49:27.343104Z"
    }
   },
   "id": "e501d8818e912d51",
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "#构建网络layers\ndef Dnnlayers(input_layers,output_layers,hidden_layers,neural):\n    layers = []\n    for i in range(hidden_layers + 2):\n        if i == 0:\n            layers.append(input_layers)\n        elif i == hidden_layers + 1:\n            layers.append(output_layers)\n        else:\n            layers.append(neural)\n    return layers\n#初值\ndef IC(x, crhoL, cuL, cpL, crhoR, cuR, cpR):\n    N = x.shape[0]\n    rho_init = np.zeros((x.shape[0], 1))\n    u_init = np.zeros((x.shape[0], 1))\n    p_init = np.zeros((x.shape[0], 1))\n\n    # rho, p - initial condition\n    for i in range(N):\n        if (x[i, 1] <= 0.0):\n            rho_init[i] = crhoL\n            u_init[i] = cuL\n            p_init[i] = cpL\n        else:\n            rho_init[i] = crhoR\n            u_init[i] = cuR\n            p_init[i] = cpR\n    U_ic = np.hstack([rho_init, u_init, p_init])\n\n    return U_ic\n#外推边界\ndef BC(x, crhoL, cuL, cpL, crhoR, cuR, cpR):\n    N = x.shape[0]\n    rho_bc = np.zeros((x.shape[0], 1))\n    u_bc = np.zeros((x.shape[0], 1))\n    p_bc = np.zeros((x.shape[0], 1))\n\n    # rho, p - initial condition\n    for i in range(N):\n        if (x[i, 1] <= 0.0):\n            rho_bc[i] = crhoL\n            u_bc[i] = cuL\n            p_bc[i] = cpL\n        else:\n            rho_bc[i] = crhoR\n            u_bc[i] = cuR\n            p_bc[i] = cpR\n    U_bc = np.hstack([rho_bc, u_bc, p_bc])\n\n    return U_bc\n",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:46.060253Z",
     "iopub.execute_input": "2024-09-24T08:23:46.061124Z",
     "iopub.status.idle": "2024-09-24T08:23:46.074353Z",
     "shell.execute_reply.started": "2024-09-24T08:23:46.061075Z",
     "shell.execute_reply": "2024-09-24T08:23:46.073511Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T08:49:27.372506Z",
     "start_time": "2024-10-06T08:49:27.358373Z"
    }
   },
   "id": "79e6a971aae89fdf",
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "class DNN(nn.Module):\n    def __init__(self, layers):\n        #继承父类\n        super(DNN, self).__init__()\n\n        #depth of network\n        self.depth = len(layers) - 1\n        #print(self.depth)\n\n        #activation of network\n        self.activation = nn.Tanh()\n\n        #create the neural network\n        layers_list = list()\n        for i in range(self.depth - 1):\n            layers_list.append(\n                ('layer_%d' % i, nn.Linear(layers[i], layers[i+1]))#create each network\n            )\n            layers_list.append(\n                ('activation_%d' % i, self.activation)  #create each activation of network\n            )\n        layers_list.append(\n            ('layer_%d' % (self.depth - 1), nn.Linear(layers[-2], layers[-1]))\n            #last network do not have activation\n        )\n        #创建一个有序字典，其中包含了从 layers_list 中获得的键值对，这在需要保持元素顺序的场景（如神经网络层的顺序）中非常有用。\n        layerDict = OrderedDict(layers_list)\n\n        #deploy layers\n        self.layers = nn.Sequential(layerDict)\n        #print(self.layers)\n\n    #forword network,output the result of network\n    def forword(self, x):\n        out = self.layers(x)\n        return out",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:46.075428Z",
     "iopub.execute_input": "2024-09-24T08:23:46.075700Z",
     "iopub.status.idle": "2024-09-24T08:23:46.086317Z",
     "shell.execute_reply.started": "2024-09-24T08:23:46.075671Z",
     "shell.execute_reply": "2024-09-24T08:23:46.085503Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T08:49:27.387546Z",
     "start_time": "2024-10-06T08:49:27.375515Z"
    }
   },
   "id": "2b5f4f99b1b8181f",
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "class Pinns():\n",
    "    def __init__(self, layers):\n",
    "        self.dnn = DNN(layers).double().to(device)    \n",
    "    #forword network,output the result of network\n",
    "    def forword(self, x):\n",
    "        y = self.dnn.forword(x)\n",
    "        rho, u, p = y[:, 0:1], y[:, 1:2], y[:, 2:]\n",
    "        return rho, u, p\n",
    "    #function loss,残差点训练\n",
    "    def loss_pde(self, x_int):\n",
    "        y = self.dnn.forword(x_int)\n",
    "        rho, u, p = y[:, 0:1], y[:, 1:2], y[:, 2:]\n",
    "        #gamma = self.gamma\n",
    "\n",
    "        U2 = rho * u\n",
    "        U3 = 0.5 * rho * u ** 2 + p / (1.4 - 1.0)#0.4 #gamma - 1 = 0.4\n",
    "        # F1 = U2\n",
    "        F2 = rho * u ** 2 + p\n",
    "        F3 = u * (U3 + p)\n",
    "\n",
    "        # Gradients and partial derivatives\n",
    "        drho_g = gradients(rho, x_int)[0]\n",
    "        rho_t, rho_x = drho_g[:, :1], drho_g[:, 1:]\n",
    "\n",
    "        dU2_g = gradients(U2, x_int)[0]\n",
    "        U2_t, U2_x = dU2_g[:, :1], dU2_g[:, 1:]\n",
    "        dU3_g = gradients(U3, x_int)[0]\n",
    "        U3_t, U3_x = dU3_g[:, :1], dU3_g[:, 1:]\n",
    "        dF2_g = gradients(F2, x_int)[0]\n",
    "        F2_t, F2_x = dF2_g[:, :1], dF2_g[:, 1:]\n",
    "        dF3_g = gradients(F3, x_int)[0]\n",
    "        F3_t, F3_x = dF3_g[:, :1], dF3_g[:, 1:]\n",
    "        \n",
    "        f = (((rho_t + U2_x) ) ** 2).mean() + \\\n",
    "            (((U2_t + F2_x) ) ** 2).mean() + \\\n",
    "            (((U3_t + F3_x) ) ** 2).mean()  #\n",
    "        \n",
    "        return f\n",
    "    #initial loss function\n",
    "    def loss_ic(self, x_ic, U_ic):\n",
    "        y_ic = self.dnn.forword(x_ic)\n",
    "        rho_ic_nn, u_ic_nn, p_ic_nn = y_ic[:, 0:1], y_ic[:, 1:2], y_ic[:, 2:]\n",
    "        loss_ics = ((rho_ic_nn - U_ic[:,0:1]) ** 2).mean()+\\\n",
    "                   ((u_ic_nn - U_ic[:,1:2]) ** 2).mean()+\\\n",
    "                   ((p_ic_nn - U_ic[:,2:]) ** 2).mean()\n",
    "        return loss_ics"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:46.087638Z",
     "iopub.execute_input": "2024-09-24T08:23:46.088500Z",
     "iopub.status.idle": "2024-09-24T08:23:46.103784Z",
     "shell.execute_reply.started": "2024-09-24T08:23:46.088466Z",
     "shell.execute_reply": "2024-09-24T08:23:46.102979Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T08:49:47.401306Z",
     "start_time": "2024-10-06T08:49:47.385650Z"
    }
   },
   "id": "5a835b6157600a38",
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "class conbine():\n",
    "    def __init__(self, layers1, layers2, layers3, layers4):\n",
    "        self.model1 = Pinns(layers1)\n",
    "        self.model2 = Pinns(layers2)\n",
    "        self.model3 = DNN(layers3).double().to(device)    #x(t)\n",
    "        self.model4 = DNN(layers4).double().to(device)#f(x,t)\n",
    "    def loss_pde(self, x_int):\n",
    "        t = x_int[:, 0:1]\n",
    "        x = x_int[:, 1:2]\n",
    "        x_predict = self.model3.forword(t)[:,0:1]\n",
    "        '''xt_predict = to_numpy(torch.hstack([t,x_predict]))\n",
    "        xt_predict = torch.tensor(xt_predict, device=device, dtype=dtype, requires_grad=True)\n",
    "        f_predict = self.model4.forword(xt_predict)'''\n",
    "        x_predict = to_numpy(x_predict)\n",
    "        x_predict = torch.tensor(x_predict, device=device, dtype=dtype, requires_grad=True)\n",
    "        f_predict = self.model4.forword(x_predict)\n",
    "        f_test = self.model4.forword(x)\n",
    "        \n",
    "        index1 = [i for i in range(len(x_int)) if f_test[i] <= f_predict[i]]\n",
    "        index2 = [i for i in range(len(x_int)) if f_test[i] > f_predict[i]]\n",
    "        x1_int = x_int[index1]\n",
    "        x2_int = x_int[index2]\n",
    "        #loss_pde(self, x_int):\n",
    "        loss_pde1 = self.model1.loss_pde(x1_int)\n",
    "        loss_pde2 = self.model2.loss_pde(x2_int)\n",
    "        loss_pde = loss_pde1 + loss_pde2\n",
    "        return loss_pde\n",
    "    #loss_ic(self, x_ic, U_ic):\n",
    "    def loss_ic(self, x_ic, U_ic):\n",
    "        t = x_ic[:, 0:1]\n",
    "        x_predict = self.model3.forword(t)[:,0:1]\n",
    "        x_predict = to_numpy(x_predict)\n",
    "        x_predict = torch.tensor(x_predict, device=device, dtype=dtype, requires_grad=True)\n",
    "        f_predict = self.model4.forword(x_predict)\n",
    "        f_test = self.model4.forword(x_ic[:, 1:2])\n",
    "        \n",
    "        index1 = [i for i in range(len(x_ic)) if f_test[i] <= f_predict[i]]\n",
    "        index2 = [i for i in range(len(x_ic)) if f_test[i] > f_predict[i]]        \n",
    "        x1_ic = x_ic[index1]\n",
    "        x2_ic = x_ic[index2]\n",
    "        U1_ic = U_ic[index1]\n",
    "        U2_ic = U_ic[index2]\n",
    "        loss_ic1 = self.model1.loss_ic(x1_ic, U1_ic)\n",
    "        loss_ic2 = self.model2.loss_ic(x2_ic, U2_ic)\n",
    "        loss_ic = loss_ic1 + loss_ic2\n",
    "        return loss_ic\n",
    "    def loss_xic(self, t_ic, x_ic):\n",
    "        x_predict = self.model3.forword(t_ic)[:,0:1]\n",
    "        loss_ic = ((x_predict - x_ic[:, 0:1]) ** 2).mean()\n",
    "        return loss_ic\n",
    "    '''def loss_f(self, t_int):\n",
    "        x_predict = self.model3.forword(t_int)[:,0:1]\n",
    "        xt_int = to_numpy(torch.hstack([t_int, x_predict]))\n",
    "        xt_int = torch.tensor(xt_int, dtype=dtype, device=device, requires_grad=True)\n",
    "        f_predict = self.model4.forword(xt_int)\n",
    "        loss_f = ((f_predict) ** 2).mean()\n",
    "        return loss_f'''\n",
    "    #R-H，D=u关系\n",
    "    def loss_rh(self, t, x_start, x_end):\n",
    "        x_predict = self.model3.forword(t)[:,0:1]\n",
    "        dx = gradients(x_predict, t)[0]\n",
    "        dx_t = dx[:, 0:1]#s_x\n",
    "        x_predict = to_numpy(x_predict)\n",
    "        x_predict = torch.tensor(x_predict, device=device, dtype=dtype, requires_grad=True)\n",
    "        f_predict = self.model4.forword(x_predict)\n",
    "        df_predict = gradients(f_predict, x_predict)[0]\n",
    "        df_x = df_predict[:, 0:1]#n_x\n",
    "        \n",
    "        index = torch.where((x_predict<=x_end)&(x_predict>=x_start))\n",
    "        t = t[index]\n",
    "        x_predict = x_predict[index]\n",
    "        dx_t = dx_t[index]\n",
    "        #\n",
    "        x_int = to_numpy(torch.cat([t.reshape(-1,1), x_predict.reshape(-1,1)], dim=1))\n",
    "        x_int = torch.tensor(x_int, dtype=dtype, requires_grad=True, device=device)\n",
    "        rho1, u1, p1 = self.model1.forword(x_int)\n",
    "        rho2, u2, p2 = self.model2.forword(x_int)\n",
    "\n",
    "        print('df_x[2]',df_x[2],'dx_t[2]',dx_t[2])\n",
    "        \n",
    "        loss_rh = ((u1 - u2) ** 2).mean()+ ((p1 - p2) ** 2).mean()+((dx_t - u1*df_x) ** 2).mean()+((dx_t - u2*df_x) ** 2).mean()\n",
    "        return loss_rh\n",
    "       "
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:46.105022Z",
     "iopub.execute_input": "2024-09-24T08:23:46.105460Z",
     "iopub.status.idle": "2024-09-24T08:23:46.123419Z",
     "shell.execute_reply.started": "2024-09-24T08:23:46.105418Z",
     "shell.execute_reply": "2024-09-24T08:23:46.122561Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T13:39:29.165429Z",
     "start_time": "2024-10-06T13:39:29.149876Z"
    }
   },
   "id": "6be409ae6f6ebf92",
   "execution_count": 48,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "#网格构建\nnum_x = 401\nnum_t = 201\nTstart = 0.0\nTend = 0.5\nXstart = -1.0\nXend = 1.0\nx = np.linspace(Xstart, Xend, num_x)\nt = np.linspace(Tstart, Tend, num_t)\nt_grid, x_grid = np.meshgrid(t, x)\nT = t_grid.flatten()[:, None]\nX = x_grid.flatten()[:, None]\nX_star = np.hstack((T.flatten()[:, None], X.flatten()[:, None]))\n# 时空边界\nlb = X_star.min(0)  # 时空下界，\nub = X_star.max(0)  # 时空上界\nprint('lb is :',lb, ', ub is :', ub)",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:46.413829Z",
     "iopub.execute_input": "2024-09-24T08:23:46.414437Z",
     "iopub.status.idle": "2024-09-24T08:23:46.432210Z",
     "shell.execute_reply.started": "2024-09-24T08:23:46.414400Z",
     "shell.execute_reply": "2024-09-24T08:23:46.431284Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T13:39:29.498224Z",
     "start_time": "2024-10-06T13:39:29.485363Z"
    }
   },
   "id": "84aca8ea777720d8",
   "execution_count": 49,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "#初值\n#loss_ic(self, x_ic, U_ic):\nnum_i_train = 201\nid_ic = np.random.choice(num_x, num_i_train, replace=False)\nirhoL, iuL, ipL = 1.0, 1.0, 0.125\nirhoR, iuR, ipR = 0.5, 1.0, 0.125\nrho_u_p_LR = [irhoL, iuL, ipL, irhoR, iuR, ipR]\nx_ic = x_grid[id_ic, 0][:, None]\nt_ic = t_grid[id_ic, 0][:, None]\nx_ic_train = np.hstack((t_ic, x_ic))\nU_ic = IC(x_ic_train, irhoL, iuL, ipL, irhoR, iuR, ipR)\n#tensor\nx_ic_train = torch.tensor(x_ic_train, dtype=dtype, requires_grad=True, device=device)\nU_ic = torch.tensor(U_ic, dtype=dtype, device=device)\n#loss_xic(self, t_ic, x_ic):\nt_ic_train = t_grid[id_ic, 0][:, None]\nx_ic = t_ic_train*1\n#tensor\nt_ic_train = torch.tensor(t_ic_train, dtype=dtype, device=device, requires_grad=True)\nx_ic = torch.tensor(x_ic, dtype=dtype, device=device)",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:46.980811Z",
     "iopub.execute_input": "2024-09-24T08:23:46.981641Z",
     "iopub.status.idle": "2024-09-24T08:23:47.183851Z",
     "shell.execute_reply.started": "2024-09-24T08:23:46.981604Z",
     "shell.execute_reply": "2024-09-24T08:23:47.183052Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T13:39:29.760753Z",
     "start_time": "2024-10-06T13:39:29.740995Z"
    }
   },
   "id": "bb1f3322bfae61ab",
   "execution_count": 50,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "#残差\n#loss_pde(self, x_int):\nnum_f_train = 5000\nid_f = np.random.choice(num_x * num_t, num_f_train, replace=False)\nx_int = X[:, 0][id_f, None]\nt_int = T[:, 0][id_f, None]\nx_int_train = np.hstack((t_int, x_int))\n#tensor\nx_int_train = torch.tensor(x_int_train, dtype=dtype, requires_grad=True, device=device)",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:47.599759Z",
     "iopub.execute_input": "2024-09-24T08:23:47.600812Z",
     "iopub.status.idle": "2024-09-24T08:23:47.609794Z",
     "shell.execute_reply.started": "2024-09-24T08:23:47.600765Z",
     "shell.execute_reply": "2024-09-24T08:23:47.608821Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T13:39:29.976736Z",
     "start_time": "2024-10-06T13:39:29.967704Z"
    }
   },
   "id": "b5ad1a076b1101fa",
   "execution_count": 51,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "#rh关系\n#loss_rh(self, t, x_start, x_end)\nt_all = t_grid[0, :][:, None]\nx_start = Xstart\nx_end = Xend\n#tensor\nt_all = torch.tensor(t_all, dtype=dtype, requires_grad=True, device=device)\nx_start = torch.tensor(x_start, dtype=dtype, device=device)\nx_end = torch.tensor(x_end, dtype=dtype, device=device)",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:48.313405Z",
     "iopub.execute_input": "2024-09-24T08:23:48.314162Z",
     "iopub.status.idle": "2024-09-24T08:23:48.321599Z",
     "shell.execute_reply.started": "2024-09-24T08:23:48.314120Z",
     "shell.execute_reply": "2024-09-24T08:23:48.320665Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T13:39:30.225682Z",
     "start_time": "2024-10-06T13:39:30.210516Z"
    }
   },
   "id": "b889c0c6fa2ecdf9",
   "execution_count": 52,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "#实例化神经网络\n",
    "layers1 = Dnnlayers(2, 3, 5, 30)\n",
    "layers2 = Dnnlayers(2, 3, 5, 30)\n",
    "layers3 = Dnnlayers(1, 1, 2, 10)\n",
    "layers4 = Dnnlayers(1, 1, 2, 10)\n",
    "# Pinns(layers)\n",
    "model = conbine(layers1, layers2, layers3, layers4)\n",
    "loss_total_history = []\n",
    "loss_pde_history = []\n",
    "loss_ibc_history = []\n",
    "loss_rh_history = []\n",
    "lr = 0.001\n",
    "optimizer = torch.optim.Adam(list(model.model1.dnn.parameters())+list(model.model2.dnn.parameters())+list(model.model3.parameters())+list(model.model4.parameters()), lr=lr)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:49.532299Z",
     "iopub.execute_input": "2024-09-24T08:23:49.533197Z",
     "iopub.status.idle": "2024-09-24T08:23:50.456244Z",
     "shell.execute_reply.started": "2024-09-24T08:23:49.533152Z",
     "shell.execute_reply": "2024-09-24T08:23:50.455427Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T13:39:30.707587Z",
     "start_time": "2024-10-06T13:39:30.686553Z"
    }
   },
   "id": "70089af1370df782",
   "execution_count": 53,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "#training\n",
    "def train(epoch):\n",
    "    model.model1.dnn.train()\n",
    "    model.model2.dnn.train()\n",
    "    model.model3.train()\n",
    "    model.model4.train()\n",
    "    def closure():\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        loss_pde = model.loss_pde(x_int_train)\n",
    "        loss_ic = model.loss_ic(x_ic_train, U_ic)\n",
    "        loss_xic = model.loss_xic(t_ic_train, x_ic)\n",
    "        #loss_f = model.loss_f(t_all)\n",
    "        loss_rh = model.loss_rh(t_all, x_start, x_end)\n",
    "                   \n",
    "        loss = loss_pde + 10*(loss_ic + loss_rh + loss_xic)#+ loss_f \n",
    "        if epoch%1==0:\n",
    "            print(f'epoch:{epoch}, loss:{loss:.8f},loss_pde:{loss_pde:.8f}, loss_IC:{loss_ic:.8f}, loss_xic:{loss_xic:.8f},loss_rh:{loss_rh:.8f}')#, loss_f:{loss_f:.8f}\n",
    "        \n",
    "        loss_total_history.append(to_numpy(loss))\n",
    "        loss_pde_history.append(to_numpy(loss_pde))\n",
    "        loss_ibc_history.append(to_numpy(loss_ic+loss_xic))\n",
    "        loss_rh_history.append(to_numpy(loss_rh))\n",
    "        \n",
    "        loss.backward()\n",
    "        return loss\n",
    "    loss = optimizer.step(closure)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:50.457846Z",
     "iopub.execute_input": "2024-09-24T08:23:50.458658Z",
     "iopub.status.idle": "2024-09-24T08:23:50.466498Z",
     "shell.execute_reply.started": "2024-09-24T08:23:50.458610Z",
     "shell.execute_reply": "2024-09-24T08:23:50.465527Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T13:39:31.098287Z",
     "start_time": "2024-10-06T13:39:31.082112Z"
    }
   },
   "id": "f674917808bcff0d",
   "execution_count": 54,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "epochs1 = 10000\n#for training\ntic = time.time()\nfor epoch in range(1, epochs1+1):\n    train(epoch)\ntoc = time.time()",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T08:23:52.150100Z",
     "iopub.execute_input": "2024-09-24T08:23:52.151126Z",
     "iopub.status.idle": "2024-09-24T09:31:41.292026Z",
     "shell.execute_reply.started": "2024-09-24T08:23:52.151082Z",
     "shell.execute_reply": "2024-09-24T09:31:41.291134Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T13:39:37.310260Z",
     "start_time": "2024-10-06T13:39:31.554079Z"
    }
   },
   "id": "59e6753e762b2257",
   "execution_count": 55,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "optimizer = torch.optim.LBFGS(list(model.model1.dnn.parameters())+list(model.model2.dnn.parameters())+list(model.model3.parameters()),lr=1.0,max_iter=10)\nepochs2 = 1000\ntic = time.time()\nfor epoch in range(1, epochs2+1):\n    train(epoch)\ntoc = time.time()",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T09:31:41.293728Z",
     "iopub.execute_input": "2024-09-24T09:31:41.294028Z",
     "iopub.status.idle": "2024-09-24T09:38:28.107433Z",
     "shell.execute_reply.started": "2024-09-24T09:31:41.293996Z",
     "shell.execute_reply": "2024-09-24T09:38:28.106494Z"
    },
    "trusted": true
   },
   "id": "cdd6739e70b2ed04",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "#test\n",
    "t_test = t_grid[0, :][:, None]\n",
    "t_test = torch.tensor(t_test, dtype=dtype, requires_grad=True, device=device)\n",
    "#x(t)\n",
    "x_predict = to_numpy(model.model3.forword(t_test)[:,0:1])\n",
    "#f(x,t)\n",
    "xt_test = np.hstack([to_numpy(t_test), x_predict])\n",
    "xt_test = torch.tensor(xt_test, dtype=dtype, requires_grad=True, device=device)\n",
    "f_predict = model.model4.forword(xt_test)\n",
    "f_exact = np.zeros_like(to_numpy(t_test))\n",
    "x_exact = to_numpy(t_test)\n",
    "#x(t)图像\n",
    "plt.figure()\n",
    "plt.plot(x_predict, to_numpy(t_test), 'r', label='x_predict')\n",
    "plt.plot(x_exact, to_numpy(t_test), 'b', label='x_exact')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('t')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "error = abs(x_predict - x_exact)\n",
    "plt.figure()\n",
    "plt.plot(to_numpy(t_test), error, label='x_error')\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('error')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "#f(x,t)\n",
    "plt.figure()\n",
    "plt.plot(x_predict, to_numpy(f_predict), 'r', label='f_predict')\n",
    "plt.plot(x_predict, f_exact, 'b', label='f_exact')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('f')\n",
    "plt.legend()\n",
    "plt.show()"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T09:38:28.108779Z",
     "iopub.execute_input": "2024-09-24T09:38:28.109549Z",
     "iopub.status.idle": "2024-09-24T09:38:28.646669Z",
     "shell.execute_reply.started": "2024-09-24T09:38:28.109498Z",
     "shell.execute_reply": "2024-09-24T09:38:28.645752Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T10:09:05.636332Z",
     "start_time": "2024-10-06T10:09:05.351703Z"
    }
   },
   "id": "6f74f9ec3398306c",
   "execution_count": 28,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "x_int_test = np.hstack([T[:, 0][:, None], X[:, 0][:, None]])\n",
    "\n",
    "#exact\n",
    "rho_exact = np.zeros_like(x_int_test[:, 0:1])\n",
    "u_exact = np.zeros_like(x_int_test[:, 0:1])\n",
    "p_exact = np.zeros_like(x_int_test[:, 0:1])\n",
    "u_exact[:, 0:1] = 1.0\n",
    "p_exact[:, 0:1] = 0.125\n",
    "\n",
    "x_int_test = torch.tensor(x_int_test, dtype=dtype, requires_grad=True, device=device)\n",
    "f_test = model.model4.forword(x_int_test)\n",
    "t = to_numpy(x_int_test[:, 0:1])\n",
    "t = torch.tensor(t, dtype=dtype, requires_grad=True, device=device)\n",
    "x = x_int_test[:, 1:2]\n",
    "x_predict = model.model3.forword(t)[:,0:1]\n",
    "xt_predict = to_numpy(torch.hstack([t, x_predict]))\n",
    "xt_predict = torch.tensor(xt_predict, dtype=dtype, requires_grad=True, device=device)\n",
    "f_predict = model.model4.forword(xt_predict)\n",
    "\n",
    "index1 = [i for i in range(len(x_int_test)) if f_test[i] <= f_predict[i]]\n",
    "index2 = [i for i in range(len(x_int_test)) if f_test[i] > f_predict[i]]\n",
    "rho_exact[index1] = 1.0\n",
    "rho_exact[index2] = 0.5\n",
    "\n",
    "x1_int = x_int_test[index1]\n",
    "x2_int = x_int_test[index2]\n",
    "rho_predict1, u_predict1, p_predict1 = model.model1.forword(x1_int)\n",
    "rho_predict2, u_predict2, p_predict2 = model.model2.forword(x2_int)\n",
    "\n",
    "rho_predict = torch.vstack([rho_predict1, rho_predict2])\n",
    "rho_predict[index1] = rho_predict1\n",
    "rho_predict[index2] = rho_predict2\n",
    "u_predict = torch.vstack([u_predict1, u_predict2])\n",
    "u_predict[index1] = u_predict1\n",
    "u_predict[index2] = u_predict2\n",
    "p_predict = torch.vstack([p_predict1, p_predict2])\n",
    "p_predict[index1] = p_predict1\n",
    "p_predict[index2] = p_predict2\n",
    "\n",
    "t_grid_test = to_numpy(x_int_test[:, 0:1]).reshape(num_x, num_t)\n",
    "x_grid_test = to_numpy(x_int_test[:, 1:2]).reshape(num_x, num_t)\n",
    "rho_predict = to_numpy(rho_predict).reshape(num_x, num_t)\n",
    "u_predict = to_numpy(u_predict).reshape(num_x, num_t)\n",
    "p_predict = to_numpy(p_predict).reshape(num_x, num_t)\n",
    "rho_exact = rho_exact.reshape(num_x, num_t)\n",
    "u_exact = u_exact.reshape(num_x, num_t)\n",
    "p_exact = p_exact.reshape(num_x, num_t)\n",
    "\n",
    "error_rho = abs(rho_predict - rho_exact)\n",
    "error_u = abs(u_predict - u_exact)\n",
    "error_p = abs(p_predict - p_exact)\n",
    "#\n",
    "plt.figure()\n",
    "plt.contourf(x_grid_test, t_grid_test, rho_predict, levels=np.linspace(rho_predict.min(), rho_predict.max(), 1000), cmap='jet')\n",
    "plt.colorbar()\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('t')\n",
    "plt.title('rho_predict')\n",
    "plt.show()\n",
    "plt.figure()\n",
    "plt.contourf(x_grid_test, t_grid_test, u_predict, levels=np.linspace(u_predict.min(), u_predict.max(), 1000), cmap='jet')\n",
    "plt.colorbar()\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('t')\n",
    "plt.title('u_predict')\n",
    "plt.show()\n",
    "plt.figure()\n",
    "plt.contourf(x_grid_test, t_grid_test, p_predict, levels=np.linspace(p_predict.min(), p_predict.max(), 1000), cmap='jet')\n",
    "plt.colorbar()\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('t')\n",
    "plt.title('p_predict')\n",
    "plt.show()\n",
    "#\n",
    "plt.figure()\n",
    "plt.contourf(x_grid_test, t_grid_test, error_rho, levels=np.linspace(error_rho.min(), error_rho.max(), 1000), cmap='jet')\n",
    "plt.colorbar()\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('t')\n",
    "plt.title('error_rho')\n",
    "plt.show()\n",
    "plt.figure()\n",
    "plt.contourf(x_grid_test, t_grid_test, error_u, levels=np.linspace(error_u.min(), error_u.max(), 1000), cmap='jet')\n",
    "plt.colorbar()\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('t')\n",
    "plt.title('error_u')\n",
    "plt.show()\n",
    "plt.figure()\n",
    "plt.contourf(x_grid_test, t_grid_test, error_p, levels=np.linspace(error_p.min(), error_p.max(), 1000), cmap='jet')\n",
    "plt.colorbar()\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('t')\n",
    "plt.title('error_p')\n",
    "plt.show()\n",
    "\n",
    "print('L2 for error_rho is :', np.sqrt(((error_rho)**2).mean()))\n",
    "print('L2 for error_u is :', np.sqrt(((error_u)**2).mean()))\n",
    "print('L2 for error_p is :', np.sqrt(((error_p)**2).mean()))\n",
    "'''\n",
    "L2 for error_rho is : 0.000971365971313373\n",
    "L2 for error_u is : 0.0006760965039426095\n",
    "L2 for error_p is : 0.00010269123556697194\n",
    "'''"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-09-24T09:57:18.511691Z",
     "iopub.execute_input": "2024-09-24T09:57:18.512071Z",
     "iopub.status.idle": "2024-09-24T09:57:38.506033Z",
     "shell.execute_reply.started": "2024-09-24T09:57:18.512032Z",
     "shell.execute_reply": "2024-09-24T09:57:38.504918Z"
    },
    "trusted": true,
    "ExecuteTime": {
     "end_time": "2024-10-06T10:17:51.582119Z",
     "start_time": "2024-10-06T10:17:43.537494Z"
    }
   },
   "id": "5dae7274de0cce36",
   "execution_count": 30,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "",
   "metadata": {},
   "id": "f56e33350875c610",
   "execution_count": null,
   "outputs": []
  }
 ]
}
